#include<torch/extension.h>
#include<vector>

void add_hip(const float *d_A, const float *d_B, float *d_C, int size);

torch::Tensor add(torch::Tensor a, torch::Tensor b) {
	auto c = torch::zeros_like(a);
	add_hip(a.data_ptr<float>(), b.data_ptr<float>(), c.data_ptr<float>(), a.numel());
	return c;
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
	m.def("add", &add, "Add two tensors (HIP)");
}
